import requests
from PIL import Image
import pytesseract

try:
    from StringIO import StringIO
except ImportError:
    from io import BytesIO as StringIO


class OCRUtil:
    """
    OCR验证码识别工具类
    """

    @staticmethod
    def get_captcha_by_url(captcha_url, tesseract_cmd=None):
        """
        通过验证码图片URL识别验证码

        Args:
            captcha_url (str): 验证码图片的URL地址
            tesseract_cmd (str, optional): tesseract可执行文件路径

        Returns:
            str: 识别出的验证码文本
        """
        # 如果指定了tesseract路径，则设置
        if tesseract_cmd:
            pytesseract.pytesseract.tesseract_cmd = tesseract_cmd

        # 从URL获取验证码图片
        response = requests.get(captcha_url)
        response.raise_for_status()

        # 将图片数据转换为Image对象
        image = Image.open(StringIO(response.content))

        # 使用OCR识别验证码
        captcha_text = pytesseract.image_to_string(image)

        # 清理识别结果，去除空格和换行符
        captcha_text = captcha_text.strip().replace(' ', '').replace('\n', '')

        return captcha_text

    @staticmethod
    def get_captcha_by_image(image_path, tesseract_cmd=None):
        """
        通过本地验证码图片识别验证码

        Args:
            image_path (str): 验证码图片的本地路径
            tesseract_cmd (str, optional): tesseract可执行文件路径

        Returns:
            str: 识别出的验证码文本
        """
        # 如果指定了tesseract路径，则设置
        if tesseract_cmd:
            pytesseract.pytesseract.tesseract_cmd = tesseract_cmd

        # 打开本地图片
        image = Image.open(image_path)

        # 使用OCR识别验证码
        captcha_text = pytesseract.image_to_string(image)

        # 清理识别结果，去除空格和换行符
        captcha_text = captcha_text.strip().replace(' ', '').replace('\n', '')

        return captcha_text

    @staticmethod
    def get_captcha_by_image_object(image, tesseract_cmd=None):
        """
        通过Image对象识别验证码

        Args:
            image (PIL.Image): 验证码图片对象
            tesseract_cmd (str, optional): tesseract可执行文件路径

        Returns:
            str: 识别出的验证码文本
        """
        # 如果指定了tesseract路径，则设置
        if tesseract_cmd:
            pytesseract.pytesseract.tesseract_cmd = tesseract_cmd

        # 使用OCR识别验证码
        captcha_text = pytesseract.image_to_string(image)

        # 清理识别结果，去除空格和换行符
        captcha_text = captcha_text.strip().replace(' ', '').replace('\n', '')

        return captcha_text




if __name__ == '__main__':
    print(OCRUtil.get_captcha_by_url('http://123.57.201.42:8091/roncoo-pay-web-boss/rcCaptcha.jpg'))
